import re
import numpy as np
import sounddevice as sd
from TTS.api import TTS as CoquiTTS_API
import torch
import collections
from TTS.utils.radam import RAdam
import librosa
import logging
import os
import contextlib

# 设置日志级别
logging.getLogger("TTS").setLevel(logging.ERROR)

# 显式添加 RAdam 和其他必要对象到安全全局列表
torch.serialization.add_safe_globals([RAdam, collections.defaultdict, dict])

_step_re = re.compile(r'^(第\d+步：)\s*')
_prefix_re = re.compile(r'^\s*\d+[^\w\s]*\s*')

class Wav:
    def __init__(self, data, sample_rate):
        self.data = data
        self.sample_rate = sample_rate

class CoquiTTS:
    def __init__(self, model_name_zh="tts_models/zh-CN/baker/tacotron2-DDC-GST", model_name_en="tts_models/en/ljspeech/vits"):
        self.tts_zh = CoquiTTS_API(model_name_zh, progress_bar=False, gpu=False)
        self.tts_en = CoquiTTS_API(model_name_en, progress_bar=False, gpu=False)
        self.synthesizer_zh = self.tts_zh.synthesizer.tts_config.audio["sample_rate"]
        self.synthesizer_en = self.tts_en.synthesizer.tts_config.audio["sample_rate"]
        self._chinese_re = re.compile(r'[\u4e00-\u9fa5]')

    def _clean_zh_text(self, text: str) -> str:
        """清理中文文本，保留标点符号或空格"""
        # 保留中文字符、标点符号和空格，去除其他字符
        return re.sub(r'[^\u4e00-\u9fa5\s，。？！、]', '', text).strip()

    def _split_text(self, text: str) -> list[tuple[str, str]]:
        """分割中英文片段，并在中文后接英文时添加逗号"""
        pattern = r'[\u4e00-\u9fa5\s，。？！、]+|[^\u4e00-\u9fa5\s，。？！、]+'
        segments = re.findall(pattern, text)
        lang_segments = []
        for i, seg in enumerate(segments):
            is_chinese = bool(self._chinese_re.search(seg))
            cleaned_seg = seg.strip()
            if not cleaned_seg:
                continue
            # 如果是中文片段且后面是英文，添加逗号
            if is_chinese and i + 1 < len(segments) and not self._chinese_re.search(segments[i + 1]):
                cleaned_seg = cleaned_seg.rstrip('，。？！、') + '，'
            lang_segments.append((cleaned_seg, 'zh' if is_chinese else 'en'))
        return lang_segments

    def _normalize(self, text: str) -> list[tuple[str, str]]:
        """规范化文本，按句子分割并标记语言"""
        text = text.strip()
        m = _step_re.match(text)
        header = m.group(1) if m else ""
        if m:
            text = text[m.end():]
        text = _prefix_re.sub("", text)
        text = header + text

        # 按标点分割句子
        raw_sentences = re.split(r'(?<=[。？！?.!])\s*', text)
        segments = []
        for sentence in raw_sentences:
            sentence = sentence.strip()
            if not sentence or re.match(r'^[。？！?.!,，、\s]*$', sentence):
                continue
            if not re.search(r'[。？！?.!]$', sentence):
                sentence += "。"
            # 进一步分割句子中的中英文
            sub_segments = self._split_text(sentence)
            segments.extend(sub_segments)
        return segments

    def speak(self, text: str):
        segments = self._normalize(text)
        wavs = []

        for seg_text, lang in segments:
            if not seg_text:
                continue
            try:
                if lang == 'zh':
                    wavs.append(Wav(self.tts_zh.tts(seg_text), self.synthesizer_zh))
                    print(f"🔊 [TTS-Coqui] 中文播报: {seg_text}")
                else:
                    wavs.append(Wav(self.tts_en.tts(seg_text), self.synthesizer_en))
                    print(f"🔊 [TTS-Coqui] 英文播报: {seg_text}")

            except Exception as e:
                print(f"生成音频失败 ({seg_text}): {e}")
                continue

        if wavs:
            # 播放音频
            for wav in wavs:
                sd.play(wav.data, samplerate=wav.sample_rate)
                sd.wait()
            print("中英文混合音频已生成并播放，保存为 output_combined.wav")
        else:
            print("未生成任何音频片段。")


def main():
    # 定义测试文本
    text = "锄禾日当午，汗滴禾下土。Hello world! 谁知盘中餐，粒粒皆辛苦。"

    # 实例化 CoquiTTS
    model_name_zh = "tts_models/zh-CN/baker/tacotron2-DDC-GST"
    model_name_en = "tts_models/en/ljspeech/vits"
    tts = CoquiTTS(model_name_zh, model_name_en)

    # 调用 speak 方法处理文本
    tts.speak(text)

if __name__ == "__main__":
    main()